# v7--2021.11.07 refine all grasps for objects in eval split
# Here, we use a simple optimization method to refine the grasp pose.
import torch
import numpy as np
import os
import sys
import configargparse

sys.path.append('./')
sys.path.append('./dif')
from grasp_refine.refine_module import PandaRefine, get_best_results_from_pth

def get_args():
    p = configargparse.ArgumentParser()
    p.add_argument('--category', type=str, default='')
    p.add_argument('--pkl_root', type=str, default='')
    p.add_argument('--lr', type=float, default=1e-3, help='learning rate')
    p.add_argument('--steps', type=int, default=10, help='refine optim steps')
    opt = p.parse_args()
    return opt


def refine_grasp(grasp_info, model, refine_pth_dir, lr=1e-3, steps=10):

    pth_dir = os.path.join(refine_pth_dir, 'refine.pth')

    obj_scale = torch.from_numpy(np.array(grasp_info['pred_s'])).float().cuda()
    shape_code = grasp_info['code']
    shape_code = torch.from_numpy(shape_code).cuda() 

    model_refine = PandaRefine(grasp_info, obj_scale, shape_code)
    model_refine.cuda()
    optimizer = torch.optim.Adam(model_refine.parameters(), lr=lr)
    loss_min = 1e10

    global_best_idx = -1
    for step in range(steps):
        optimizer.zero_grad()
        losses = model_refine(model)
        loss_refine = losses['loss']
        loss_choose = losses['loss_choose']
        best_idx = losses['best_idx'] 
        loss_refine.backward(retain_graph=True)
        optimizer.step()
        
        # Save the best model
        if loss_choose < loss_min:
            loss_min = loss_choose
            global_best_idx = best_idx
            torch.save(model_refine.state_dict(), pth_dir)
    deltas = torch.load(pth_dir)
    grasp_info = get_best_results_from_pth(grasp_info, deltas, global_best_idx)
    return grasp_info